# SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC.
# SPDX-License-Identifier: Apache-2.0


import os
from pathlib import Path
from typing import Any

import pytest
import torch
import torch.nn as nn
from loguru import logger

import ttnn
from models.demos.deepseek_v3.tt.ccl import CCL
from models.demos.deepseek_v3.tt.lm_head1d import LMHead1D
from models.demos.deepseek_v3.utils.config_helpers import _check_weights_exist_and_convert, sub_state_dict
from models.demos.deepseek_v3.utils.run_config import create_run_config
from models.demos.deepseek_v3.utils.test_utils import assert_hidden_dim_pcc, get_model_config, run_module_forward


class DeepseekV3LMHead(nn.Module):
    """
    Language model head for Deepseek V3.
    """

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

    def forward(self, hidden_states):
        return self.lm_head(hidden_states)


@pytest.mark.parametrize(
    "device_params",
    [
        {"fabric_config": ttnn.FabricConfig.FABRIC_1D},
    ],
    indirect=True,
)
@pytest.mark.parametrize(
    "mode, batch_size_per_row",
    [
        ("decode", 32),
        ("prefill", 1024),
    ],
)
def test_forward_pass(
    mode: str,
    batch_size_per_row: int,
    hf_config: Any,
    mesh_device: ttnn.Device,
    ccl: CCL,
    cache_path: Path,
    set_deterministic_env: Any,
):
    reference_model = DeepseekV3LMHead(hf_config).eval()
    state_dict = sub_state_dict(reference_model.state_dict(), "lm_head.")
    batch_size = batch_size_per_row * mesh_device.shape[0]
    torch_input = torch.randn(1, 1, batch_size, hf_config.hidden_size)
    reference_output = reference_model(torch_input)

    # Setup: Convert weights and get weight_config
    weight_cache_path = (
        cache_path
        / "tests_cache"
        / os.environ.get("PYTEST_CURRENT_TEST")
        / f"{hf_config.num_hidden_layers}_layers"
        / f"mesh_{mesh_device.shape[0]}x{mesh_device.shape[1]}"
    )

    weight_config = LMHead1D.convert_weights(hf_config, (state_dict,), weight_cache_path, mesh_device)
    _check_weights_exist_and_convert(weight_cache_path, weight_config)
    model_config = get_model_config(LMHead1D, mode, mesh_device)
    model_state = LMHead1D.create_state(mesh_device, ccl)
    run_config = create_run_config(model_config, weight_config, model_state)

    # Convert input to TTNN
    tt_input = ttnn.from_torch(
        torch_input,
        device=mesh_device,
        mesh_mapper=ttnn.ShardTensor2dMesh(mesh_device, dims=(-2, None), mesh_shape=mesh_device.shape),
        dtype=ttnn.bfloat16,
        memory_config=run_config["input_memory_config"],
        layout=ttnn.TILE_LAYOUT,
    )

    # TTNN forward pass
    tt_input = ttnn.to_memory_config(tt_input, run_config["input_memory_config"])
    tt_output = run_module_forward(LMHead1D, mode, tt_input, run_config)

    expected_output_memory_config = run_config["output_memory_config"]

    # Verify output memory config matches expected
    actual_output_memory_config = tt_output.memory_config()
    assert (
        actual_output_memory_config == expected_output_memory_config
    ), f"Output memory config mismatch: expected {expected_output_memory_config}, got {actual_output_memory_config}"

    logger.info("running ttnn.to_torch")
    tt_output_torch = ttnn.to_torch(
        tt_output, mesh_composer=ttnn.ConcatMesh2dToTensor(mesh_device, dims=(-2, -1), mesh_shape=mesh_device.shape)
    )
    logger.info(f"tt_output_torch shape: {tt_output_torch.shape}")
    logger.info("finished ttnn.to_torch")
    assert_hidden_dim_pcc(tt_output_torch, reference_output, pcc_required=0.98)
